# -*- coding:utf8 -*-
"""
将句子以及中间提取的内容转换为bioe的格式
"""
from __future__ import print_function, unicode_literals
from collections import defaultdict
from functools import partial
from logger_utility.loggconfig import getLogger

logger = getLogger("Word2IEO")


class StructedTaggerSBIE(object):
    """
    针对具体的对象进行结构化的标签
    """
    SYMBOL_S = 'S-'
    SYMBOL_B = 'B-'
    SYMBOL_I = 'I-'
    SYMBOL_E = 'E-'
    SYMBOL_O = 'O'

    def extract_to_tag_with_pos(self, words, extracts, extract_max_len=None):
        """
        将句子转换为tag标注序列
        :param words:　分词后的词语序列
        :param extracts:　{'tag-base': {'刘德华': [8, ...]}, ...} 识别出来的词组以及索引位置
        :param extract_max_len:
        :return:
        """
        if extract_max_len is None:
            extract_max_len = max(len(j) for i in extracts.values() for j in i)
        words_len = len(words)
        y_tags = []

        i = 0
        find = set()
        while i < words_len:
            for len_ in range(extract_max_len, -1, -1):
                if i + len_ > words_len:
                    continue

                found = False
                str_ = ''.join(words[i: i + len_])
                for tag_base, strs in extracts.items():
                    postions = strs.get(str_, None)
                    if postions is None:
                        continue
                    if postions and i not in postions:
                        continue

                    found = True
                    break

                if found:
                    find.add(str_)
                    if len_ == 1:
                        y_tags.append(self.SYMBOL_S + tag_base)
                    else:
                        y_tags.append(self.SYMBOL_B + tag_base)
                        for _ in range(len_ - 2):
                            y_tags.append(self.SYMBOL_I + tag_base)
                        y_tags.append(self.SYMBOL_E + tag_base)
                    i += len_
                    break
            else:
                i += 1
                y_tags.append(self.SYMBOL_O)

        return y_tags, find

    def extract_to_tag_no_pos(self, words, extracts, extract_max_len=None):
        """
        将句子转换为tag标注序列
        :param words:　分词后的词语序列
        :param extracts:　{'tag-base': ['刘德华',...], ...} 识别出来的词组以及索引位置
        :param extract_max_len:
        :return:
        """
        if extract_max_len is None:
            try:
                extract_max_len = max(len(j) for i in extracts.values() for j in i)
            except:
                print(extracts)
                extract_max_len = 5
        words_len = len(words)
        y_tags = []

        i = 0
        find = set()
        while i < words_len:
            for len_ in range(extract_max_len, -1, -1):
                if i + len_ > words_len:
                    continue

                found = False
                str_ = ''.join(words[i: i + len_])
                for tag_base, strs in extracts.items():
                    if str_ not in strs:
                        continue

                    found = True
                    break

                if found:
                    find.add(str_)
                    if len_ == 1:
                        y_tags.append(self.SYMBOL_S + tag_base)
                    else:
                        y_tags.append(self.SYMBOL_B + tag_base)
                        for _ in range(len_ - 2):
                            y_tags.append(self.SYMBOL_I + tag_base)
                        y_tags.append(self.SYMBOL_E + tag_base)
                    i += len_
                    break
            else:
                i += 1
                y_tags.append(self.SYMBOL_O)

        return y_tags, find

    def tag_to_extract(self, words, y_tags, with_pos=True):
        """
        将tag转换为实际的短语
        :param words:
        :param y_tags:
        :return:
        """
        extract_list = defaultdict(partial(defaultdict, list))
        word_list = []
        begin_index = 0

        for index, (word, y_tag) in enumerate(zip(words, y_tags)):
            if y_tag.startswith('E-') or y_tag.startswith('O') or y_tag.startswith('S-'):
                if y_tag.startswith('E-'):
                    if not word_list:
                        begin_index = index
                    word_list.append(word)
                    type_, tag_base = y_tag.split('-')

                if word_list:
                    extract_list[tag_base][''.join(word_list)].append(begin_index)
                    word_list = []

                if y_tag.startswith('S-'):
                    type_, tag_base = y_tag.split('-')
                    extract_list[tag_base][word].append(index)

            else:
                if '-' in y_tag:
                    type_, tag_base = y_tag.split('-')
                    if not word_list:
                        begin_index = index
                    word_list.append(word)

        if word_list:
            extract_list[tag_base][''.join(word_list)].append(begin_index)

        if not with_pos:
            return {
                tag_base: list(item_pos.keys()) for tag_base, item_pos in extract_list.items()
            }

        return extract_list


class StructedTaggerIE(StructedTaggerSBIE):
    """
    只采用IE两种结构进行标注
    """
    SYMBOL_B = 'I-'
    SYMBOL_S = 'E-'


if __name__ == '__main__':
    def extract_to_tag_test():
        words = '2018年|3月|16日|刘|德华|去|内蒙古|呼和浩特市|广播|大学|参加|中国|人民|共和国|组织|的|音乐|颁奖|大会|，|王杰|亦|到场|出席'.split("|")
        y_tag, find = StructedTaggerSBIE().extract_to_tag(
            words,
            {
                'PER': {"王杰": [], "刘德华": []},
                'TIM': {"2018年3月16日": []},
                'LOC': {'内蒙古呼和浩特市广播大学': []},
                'ORG': {'中国人民共和国': []},
                'COF': {'音乐颁奖大会': []}
            }
        )
        for index, (word, tag) in enumerate(zip(words, y_tag)):
            print(index, word, tag)
        return words, y_tag

    def tag_to_extract_test():
        words, tags = extract_to_tag_test()
        extract_list = StructedTaggerSBIE().tag_to_extract(words, tags)
        import json
        print(json.dumps(extract_list, ensure_ascii=False, indent=2))

    tag_to_extract_test()


